import os
import torch
import torch.optim as optim
from torch.utils.data import DataLoader
from tqdm import tqdm
from data import ImageDataset
from util import linear_beta_schedule , p_losses
from Unet import SimpleUNet
from sample import sample
from random import randint
from Config import Config

# Training function
def train():
    #数据集建立
    dataset = ImageDataset(root_dir=Config.data_dir)
    dataloader = DataLoader(dataset, batch_size=Config.batch_size, shuffle=True)

    # Pre-calculate different terms for closed form
    betas = linear_beta_schedule(Config.timesteps)
    alphas = 1. - betas
    alphas_cumprod = torch.cumprod(alphas, axis=0)#前缀积
    alphas_cumprod_prev = torch.cat([torch.ones(1), alphas_cumprod[:-1]])#前一步前缀积
    model = SimpleUNet(Config.channels , Config.timesteps,Config.image_size).to(Config.device)
    optimizer = optim.Adam(model.parameters(), lr=Config.learning_rate)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=10, factor=0.5)
    
    best_loss = float('inf')
    
    for epoch in range(Config.epochs):
        epoch_loss = 0.0
        model.train()
        
        progress_bar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{Config.epochs}")
        for batch in progress_bar:
            optimizer.zero_grad()
            
            batch = batch.to(Config.device)
            # Random timesteps for each sample in the batch
            t = torch.randint(0, Config.timesteps, (batch.size(0),), device=Config.device).long()
            
            # Noise prediction loss
            loss = p_losses(model, batch, t, alphas_cumprod)
            
            loss.backward()
            optimizer.step()
            
            epoch_loss += loss.item()
            progress_bar.set_postfix(loss=loss.item())
        
        avg_loss = epoch_loss / len(dataloader)
        scheduler.step(avg_loss)
        print(f"Epoch {epoch+1}, Average Loss: {avg_loss:.4f}")
        
        # Save best model
        if avg_loss < best_loss:
            best_loss = avg_loss
            torch.save(model.state_dict(), os.path.join(Config.save_dir, "best_model.pth"))
            print(f"Saved new best model with loss: {best_loss:.4f}")
        if((epoch+1)%50==0 and Config.isplot):
            sample(model,Config)

if __name__ == "__main__":
    train()
